`include "mycpu.h"
module mem_stage(
    input clk,
    input reset,
    //WB stage allowin
    input ws_allowin,
    //MEM stage allowin
    output ms_allowin,
    //EXE to MEM
    input es_to_ms_valid,
    input [`ES_TO_MS_WD - 1:0] es_to_ms_bus,
    //MEM to WB
    output ms_to_ws_valid,
    output [`MS_TO_WS_WD - 1:0] ms_to_ws_bus,
    //read data from data sram
    input [31:0] data_sram_rdata,
    //to ds for depend
    output out_ms_valid
);

reg ms_valid;
wire ms_ready_go;
reg [`ES_TO_MS_WD - 1:0] es_to_ms_bus_r;

assign out_ms_valid = ms_valid;

assign ms_ready_go = 1'b1;

assign ms_allowin = !ms_valid || ms_ready_go && ws_allowin;

always @(posedge clk) begin
    if(reset)
        ms_valid <= 1'b0;
    else if(ms_allowin)
        ms_valid <= es_to_ms_valid;
end
always @(posedge clk) begin
    if(ms_allowin && es_to_ms_valid)
        es_to_ms_bus_r <= es_to_ms_bus;
end
wire [31:0] ms_pc;
wire [31:0] ms_mem_result;
wire [31:0] ms_alu_result;
wire [31:0] ms_final_result;
wire [31:0] ms_ld_result;
wire [4:0] ms_dest;
wire ms_res_from_mem;
wire ms_dst_is_r1;
wire ms_gr_we;
wire ms_op_ls_b;
wire ms_op_ls_h;
wire ms_ld_unsigned;
wire [1:0] ms_ls_vaddr;

assign {
    ms_ls_vaddr[1:0],
    ms_ld_unsigned,
    ms_op_ls_b,
    ms_op_ls_h,
    ms_pc[31:0],
    ms_alu_result[31:0],
    ms_res_from_mem,
    ms_gr_we,
    ms_dest[4:0] } = es_to_ms_bus_r[`ES_TO_MS_WD - 1:0];

assign ms_mem_result   = data_sram_rdata;
assign ms_ld_result    = {32{ms_op_ls_b & ms_ls_vaddr == 2'b00}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[7]}}, ms_mem_result[7:0]}
                       | {32{ms_op_ls_b & ms_ls_vaddr == 2'b01}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[15]}}, ms_mem_result[15:8]}
                       | {32{ms_op_ls_b & ms_ls_vaddr == 2'b10}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[23]}}, ms_mem_result[23:16]}
                       | {32{ms_op_ls_b & ms_ls_vaddr == 2'b11}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[31]}}, ms_mem_result[31:24]}
                       | {32{ms_op_ls_h & ms_ls_vaddr == 2'b00}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[15]}}, ms_mem_result[15:0]}
                       | {32{ms_op_ls_h & ms_ls_vaddr == 2'b10}} & 
                         {{24{~ms_ld_unsigned & ms_mem_result[31]}}, ms_mem_result[31:16]}
                       | {32{~ms_op_ls_b & ~ms_op_ls_h}} & 
                         ms_mem_result[31:0];
assign ms_final_result = ms_res_from_mem ? ms_ld_result : ms_alu_result;

//MEM to WB
assign ms_to_ws_valid = ms_valid && ms_ready_go;
assign ms_to_ws_bus[`MS_TO_WS_WD - 1:0] = {
    ms_pc[31:0],           //69:38         
    ms_gr_we,              //37:37
    ms_dest[4:0],          //36:32
    ms_final_result[31:0]  //31:0
};

endmodule